package cn.vastbase.jdbc.io;

import cn.com.vastbase.copy.CopyManager;
import cn.com.vastbase.core.BaseConnection;
import org.junit.Test;

import java.io.*;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.SQLException;

public class CopyTestCase {

    static String cname = "cn.com.vastbase.Driver";
    static Connection conn = null;
    static String url = "jdbc:vastbase://192.168.56.17:6000/vastbase";
    static String username = "test";
    static String passwd = "Vastdata#0730";
    static String tableName = "test_io";
    static String delimiter = "|";              //定义分隔符
    static String encoding = "UTF8";            //定义字符集
    static FileInputStream is = null;
    static FileOutputStream os = null;

    public Connection conn() {
        try {
            Class.forName(cname);
            return conn = DriverManager.getConnection(url, username, passwd);
        } catch (ClassNotFoundException | SQLException e) {
            System.out.println("[FAIL] conn database fail." + e.getMessage());
        }
        return null;
    }

    public void close() {
        try {
            if (conn != null) {
                conn.close();
            }
            if (is != null) {
                is.close();
            }
            if (os != null) {
                os.close();
            }
        } catch (SQLException e) {
            e.printStackTrace();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    /*
     从本地文件导入数据库
     数据准备
     seq 100000| awk '{print $0"|copy"}' > copy_input.txt
     create table test_io(id int, name varchar(128));
     */
    @Test
    public void copyFromFileToDb() {
        try {
            is = new FileInputStream("D://copy_input.txt");
            conn();
            conn.setAutoCommit(false);
            if (conn == null) {
                return;
            }
            BaseConnection baseConn = (BaseConnection) conn;
            String sql = "Copy " + tableName + " from STDIN DELIMITER " + "'" + delimiter + "'" + " ENCODING " + "'" + encoding + "'";
            CopyManager cp = new CopyManager(baseConn);
            // 缓冲流导入，每次1024字节
            cp.copyIn(sql, is, 1024);
            conn.commit();
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            close();
        }
    }


    @Test
    public void multiCopyFromFileToDb() throws InterruptedException {
        for (int i = 0; i < 2; i++) {
            int finalI = i;
            Thread thread = new Thread(() -> {
                Connection conn = null;
                try {
                    is = new FileInputStream("D://copy_input"+(finalI +1)+".txt");
                    Class.forName(cname);
                    conn = DriverManager.getConnection(url, username, passwd);
                    if (conn == null) {
                        return;
                    }
                    conn.setAutoCommit(false);
                    BaseConnection baseConn = (BaseConnection) conn;
                    String sql = "Copy " + tableName + " from STDIN DELIMITER " + "'" + delimiter + "'" + " ENCODING " + "'" + encoding + "'";
                    CopyManager cp = new CopyManager(baseConn);
                    cp.copyIn(sql, is, 1024);
                    conn.commit();
                } catch ( Exception e) {
                    throw new RuntimeException(e);
                }finally {
                    close();
                }
            });
            thread.start();
        }

        Thread.sleep(100000);
    }


    /**
     * 从数据库导出到本地文件
     */
    @Test
    public void copyFromDbToFile() {
        try {
            if (conn() == null) {
                return;
            }
            os = new FileOutputStream("E://copy_output.txt");
            BaseConnection baseConn = (BaseConnection) conn;
            String sql = "Copy (select * from " + tableName + ") to STDOUT  DELIMITER " + "'" + delimiter + "'" + " ENCODING " + "'" + encoding + "'";
            CopyManager cp = new CopyManager(baseConn);
            cp.copyOut(sql, os);
            baseConn.close();
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            close();
        }
    }


}
